use std::{
    sync::{mpsc, Arc, Mutex},
    thread,
};

type Job = Box<dyn FnOnce() + Send + 'static>;

enum Message {
    NewJob(Job),
    Terminate,
}

/// 简单的线程池实现
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>,
}

impl ThreadPool {
    /// 创建包含 `size` 条工作线程的线程池
    ///
    /// # Panics
    ///
    /// 当 `size == 0` 时将 panic
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0, "ThreadPool size must be greater than zero.");

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));
        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }

    /// 将任务提交到线程池
    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        if let Err(err) = self.sender.send(Message::NewJob(job)) {
            eprintln!("Failed to send job to worker: {err:?}");
        }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        println!("Sending terminate messages to all workers.");

        for _ in &self.workers {
            if let Err(err) = self.sender.send(Message::Terminate) {
                eprintln!("Failed to send terminate message: {err:?}");
            }
        }

        println!("Shutting down workers...");

        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                if let Err(err) = thread.join() {
                    eprintln!("Worker {} panicked: {err:?}", worker.id);
                }
            }
        }
    }
}

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = {
                let lock_result = receiver.lock();
                if lock_result.is_err() {
                    eprintln!("Worker {id}: failed to acquire lock; terminating.");
                    break;
                }

                match lock_result.unwrap().recv() {
                    Ok(message) => message,
                    Err(_) => {
                        eprintln!("Worker {id}: channel disconnected; terminating.");
                        break;
                    }
                }
            };

            match message {
                Message::NewJob(job) => {
                    println!("Worker {id} got a job; executing.");
                    job();
                }
                Message::Terminate => {
                    println!("Worker {id} received terminate signal.");
                    break;
                }
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}